import cv2
import numpy as np
from scipy.optimize import curve_fit
from matplotlib import pyplot as plt
from skimage.feature import blob_log
def in_interval(a, l, r):
    if (a >= l) and (a <= r):
        return True
    else:
        return False


def search_light_blocks(bin_data):
    n, m = bin_data.shape
    flag_grid = np.zeros((n, m), dtype='int')
    light_blocks = []
    for i in range(n):
        for j in range(m):
            if (not flag_grid[i, j]) and bin_data[i, j]:
                queue = [(i, j)]
                l = 0
                flag_grid[i, j] = 1
                x, y = i, j
                while l < len(queue):
                    x, y = queue[l]
                    for dx in range(-1, 2):
                        for dy in range(-1, 2):
                            if (dx, dy) != (0, 0):
                                x2 = x + dx
                                y2 = y + dy
                                if in_interval(x2, 0, n - 1) and in_interval(y2, 0, m - 1):
                                    if (not flag_grid[x2, y2]) and bin_data[x2, y2]:
                                        queue.append((x2, y2))
                                        flag_grid[x2, y2] = 1
                    l += 1

                light_blocks.append(tuple(queue))
    return light_blocks


def find_ions(data, threshold_block = 12,debug=0):
    """
    This function is used to find ions in a given image data.

    The function applies a median blur to the image and calculates a threshold for binarization.
    Then, it searches for light blocks in the binarized image, and checks each light block.
    If a block contains more than a certain number of bright pixels, it is considered an ion and its center and radius are calculated.

    Parameters:

    data: np.array
        The input image data as a numpy array. It should be either grayscale (single-channel) or BGR (three-channel).
        The data type of the image MUST be np.float32.

    threshold_block: int, optional
        The minimum number of bright pixels that should be present in a block for it to be considered an ion.
        It is used to differentiate between noise points and light spots in the image.
        The default value is 12, meaning any block with more than 12 bright pixels will be considered as an ion.
        
    Returns:

    circles: list of tuples
        Each tuple represents a detected ion, where the first two elements are the x and y coordinates of the ion's center, 
        and the third element is the calculated radius of the ion.
    """
    if debug==1:
        plt.figure(figsize=(8, 6))
        print("this is raw data")
        plt.imshow(data)
    median = cv2.medianBlur(data, 5)
    if debug==1:
        plt.figure(figsize=(8, 6))
        print("after medianblur")
        plt.imshow(median)
    median = np.array(median, dtype='int')
    n, m = median.shape
    median_1D = median.reshape(n * m)

    median_1D = np.sort(median_1D)

    valley, peak = median_1D[10], median_1D[-10]
    if debug==1:
        print("valley, peak",valley, peak)
    truncate_high = median_1D[np.where(median_1D > (valley + peak) / 2)]
    truncate_low = median_1D[np.where(median_1D <= (valley + peak) / 2)]

    hist_low, bins_low = np.histogram(a=truncate_low, bins=21)
    hist_high, bins_high = np.histogram(a=truncate_high, bins=21)

    max_index_low = np.argmax(hist_low)
    max_index_high = np.argmax(hist_high)
    if debug==1:
        print("this is low count hist")
        print(max_index_low)
        plt.figure(figsize=(8, 6))
        plt.hist(bins_low[:-1], bins_low, weights=hist_low, edgecolor="black")
        plt.show()
        print("this is high count hist")
        print(max_index_high)
        plt.figure(figsize=(8, 6))
        plt.hist(bins_high[:-1], bins_high, weights=hist_high, edgecolor="yellow")
        plt.show()
    max_low = (bins_low[max_index_low] + bins_low[max_index_low + 1]) / 2
    max_high = (bins_high[max_index_high] + bins_high[max_index_high + 1]) / 2
    threshold = (max_low + max_high) / 2
    if debug==1:
        print("max_low",max_low,"max_high",max_high)
        print("threshold",threshold)
    # print("binaryzation threshold = ", threshold)
    # print(hist_high)
    # 得到threshold之后，做二值化处理(binarization)
    bin_data = np.zeros((n, m), dtype='int')
    bin_data[np.where(median > threshold)] = 1
    if debug==1:
        print("this is bin_data")
        plt.figure(figsize=(8, 6))
        plt.imshow(bin_data)
    light_blocks = search_light_blocks(bin_data)    

    circles = []
    for block in light_blocks:
        if len(block) > threshold_block:
            # we find a circle!
            sum_xy = np.sum(np.array(block), 0)
            avrg_x = sum_xy[0] / len(block)
            avrg_y = sum_xy[1] / len(block)
            r = 0
            for x, y in block:
                if (x - avrg_x) ** 2 + (y - avrg_y) ** 2 > r ** 2:
                    r = ((x - avrg_x) ** 2 + (y - avrg_y) ** 2) ** 0.5
            if r < min(n, m) / 3:
                circles.append((avrg_x, avrg_y, r))
    if debug==1:
        print(len(circles))
        print(circles)
    return np.array(circles)

def find_ionsnew(data):
    """
    This function is used to find ions in a given image data. TODO

    The function applies a median blur to the image and calculates a threshold for binarization.
    Then, it searches for light blocks in the binarized image, and checks each light block.
    If a block contains more than a certain number of bright pixels, it is considered an ion and its center and radius are calculated.

    Parameters:

    data: np.array
        The input image data as a numpy array. It should be either grayscale (single-channel) or BGR (three-channel).
        The data type of the image MUST be np.float32.

    
        
    Returns:

    : list of tuples
        Each tuple represents a detected ion, where the first two elements are the x and y coordinates of the ion's center, 
        and the third element is the calculated radius of the ion.
    """
    image_averagedmedian=cv2.medianBlur(data, 5)
    ion_position_list =blob_log(image_averagedmedian, max_sigma=5 , threshold=1000)
    ion_position_list=np.array(ion_position_list)
    #修正blob_log函数计算的半径，与以前的代码尽可能保持一致
    ion_position_list[:,2]=ion_position_list[:,2]*1.5
    sorted_indices = np.lexsort((ion_position_list[:,1], ion_position_list[:,0]))
    ion_position_list=ion_position_list[sorted_indices]
    return np.array(ion_position_list)


def gaussian(xdata, A, x0, y0, sigmax, sigmay):
    return A * np.exp(-(xdata[0] - x0) ** 2 / (2 * sigmax ** 2) - (xdata[1] - y0) ** 2 / (2 * sigmay ** 2))


# Where: a = -1/(2sigmax^2), b = -1/(2sigmay^2), c = x0/(sigmax^2), d = y0/(sigmay^2)
#       e = ln(Amp) - x0^2/(2sigmax^2) - y0^2/(2sigmay^2)
def ln_gaussian(xdata, a, b, c, d, e):
    return a * xdata[0] * xdata[0] + b * xdata[1] * xdata[1] + c * xdata[0] + d * xdata[1] + e


def fit_gaussian(data, circles):
    circles_data = []
    n, m = data.shape
    for circle in circles:
        circle_points = []
        z = []
        x, y = circle[0], circle[1]
        radius = circle[2]
        lx = round(np.floor(x - radius))
        rx = round(np.ceil(x + radius))
        ly = round(np.floor(y - radius))
        ry = round(np.ceil(y + radius))
        for i in range(lx, rx):
            for j in range(ly, ry):
                if (i - x) ** 2 + (j - y) ** 2 <= radius ** 2:
                    if in_interval(i, 0, n - 1) and in_interval(j, 0, m - 1):
                        circle_points.append([i, j])
                        z.append(data[i, j])
        circle_points = np.array(circle_points).T
        z = np.array(z)
        circles_data.append((circle_points, z))

    amp_list = []
    for i in range(len(circles_data)):
        xdata = circles_data[i][0]
        z = circles_data[i][1]
        x, y, radius = circles[i]
        amp0 = data[round(x), round(y)]
        p0 = (amp0, x, y, 1, 1)
        popt, pcov = curve_fit(gaussian, xdata, z, p0=p0, maxfev=5000)
        # print(popt)
        # amp == popt[0]
        amp_list.append(popt[0])
    return amp_list


def fit_ln_gaussian(data, circles):
    '''
    fig = plt.figure()
    ax = Axes3D(fig)
    ax.plot_surface(x,y,z)
    '''

    '''
    path = 2
    if path == 1:
        lnz = np.log(z)
        p0 = (0,0,0,0,0)
        popt, pcov = curve_fit(ln_gaussian, data, lnz, p0 = p0)
        print(popt)
        a,b,c,d,e = popt

        sigma_x = (-1/(2*a)) ** 0.5
        sigma_y = (-1/(2*b)) ** 0.5
        x0 = c * sigma_x**2
        y0 = d * sigma_y**2
        lnAmp = e + c/2 + d/2
        Amp = np.exp(lnAmp)
        print(lnAmp, Amp, x0, y0, sigma_x, sigma_y)
    elif path == 2:
    '''
    pass
